import torch
from torch.autograd import Variable

"""
nn 模块
PyTorch:nn

计算图和 autograd 是定义复杂运算符和自动求导的的一个非常强大的范例。然而对于大规模的神经网络， 原始的 autograd 可能有点太低级了。
当构建神经网络时，我们经常想到把计算组织维层级结构，其中一些具有可学习的参数，这些参数将在学习期间被优化。
在 Tensorflow 中，像 Keras， TensorFlow-Slim 和 TFLearn 这样的软件包提供了对原始图的更高级的抽象，这对于构建神经网络很有用。
在 PyTorch 中， nn 包提供了同样的功能。 nn 包提供了一组模块，他们大致相当于神经网络层。
一个模块接收一个输入变量并计算输出向量，也可能保存内部状态，如包含可学习的参数。 nn 包还定义了一组训练神经网络时有用的损失函数。
在这个例子中我们使用 nn 包来实现我们的两层网络：
"""

# N 批的大小
# D_in 输入的数据纬度
# H 隐藏层纬度
# D_out 输出数据纬度
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# 创建随机的张量保存输入和输出，放到变量里
# Create random Tensors to hold inputs and outputs, and wrap them in Variables.
x = Variable(torch.randn(N, D_in))
y = Variable(torch.randn(N, D_out), requires_grad=False)

# 使用nn包定义我们我们的模型作为序列层，nn.Sequential是一个模块也包含其它模块
# 应用它们在序列中产生输出，每一个线性的模块使用线性函数从输入计算输出，并且为权重和偏置保存内部变量
# Use the nn package to define our model as a sequence of layers. nn.Sequential
# is a Module which contains other Modules, and applies them in sequence to
# produce its output. Each Linear Module computes output from input using a
# linear function, and holds internal Variables for its weight and bias.
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out)
)

# nn包也包含主流的损失函数，现在我们将会用均方差作为我们的损失函数
# The nn package also contains definitions of popular loss functions; in this
# case we will use Mean Squared Error (MSE) as our loss function.
loss_fn = torch.nn.MSELoss(reduction='sum')

# 学习率
learning_rate = 1e-4

# 开始学习
for t in range(500):
    # 正向传播：通过传递x到模型，计算预测的y，模块对象重写了__call__操作，因此可以
    # 像函数一样进行调用。当传递一个输入值的变量到模块，它会生成输出数据的变量
    # Forward pass: compute predicted y by passing x to the model. Module objects
    # override the __call__ operator so you can call them like functions. When
    # doing so you pass a Variable of input data to the Module and it produces
    # a Variable of output data.
    y_pred = model(x)

    # 计算和打印损失，我们传递包含预测的变量和实际的y值，损失函数返回包含损失的变量
    # Compute and print loss. We pass Variables containing the predicted and true
    # values of y, and the loss function returns a Variable containing the
    # loss.
    loss = loss_fn(y_pred, y)
    print(t, loss.data)

    # 反向传播之前清零梯度
    # Zero the gradients before running the backward pass.
    model.zero_grad()

    # 反向传播：对于模型的所有的可学习的参数计算损失的梯度，
    # 在内部，对于requires_grad=True的每个模块的参数被保存在变量里
    # 因此这样调用可以计算模型的所有可学习的参数的梯度
    # Backward pass: compute gradient of the loss with respect to all the learnable
    # parameters of the model. Internally, the parameters of each Module are stored
    # in Variables with requires_grad=True, so this call will compute gradients for
    # all learnable parameters in the model.

    loss.backward()

    # 更新参数
    for param in model.parameters():
        param.data -= learning_rate * param.grad.data